from pymilvus import connections, FieldSchema, CollectionSchema, DataType, Collection, utility
import numpy as np


def init_milvus(collection_name="insurance_manual_01"):
    """
    初始化Milvus连接并创建集合（表）

    参数:
        collection_name: 集合名称（类似数据库表名）

    返回:
        Milvus集合对象
    """
    # 连接到Milvus服务
    connections.connect(host="localhost", port="19530")

    # 检查集合是否已存在，如果存在就删除（方便测试，实际应用中谨慎使用）
    if utility.has_collection(collection_name):
        utility.drop_collection(collection_name)

    # 定义字段（类似数据库表的列）
    fields = [
        # 主键字段：自动递增的整数
        FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
        # 页码字段：存储文本所在的PDF页码
        FieldSchema(name="page_num", dtype=DataType.INT64),
        # 文本片段字段：存储原始文本
        FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=2000),  # 最大2000字符
        # 向量字段：存储1024维的向量
        FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=1024)
    ]

    # 创建集合 schema
    schema = CollectionSchema(fields, description="平安保险用户手册向量集合")

    # 创建集合
    collection = Collection(name=collection_name, schema=schema)

    # 创建向量索引（加速相似度搜索）
    index_params = {
        "index_type": "IVF_FLAT",  # 基础索引类型，适合中小规模数据
        "metric_type": "L2",  # 距离计算方式：L2欧式距离（值越小越相似）
        "params": {"nlist": 128}  # 索引参数，128是推荐值
    }
    collection.create_index(field_name="embedding", index_params=index_params)

    print(f"Milvus集合 '{collection_name}' 创建成功")
    return collection


def insert_into_milvus(collection, chunks, embeddings):
    """
    将文本片段和对应的向量插入到Milvus集合

    参数:
        collection: Milvus集合对象
        chunks: 文本片段列表（包含page_num和chunk）
        embeddings: 向量列表（与文本片段一一对应）
    """
    # 准备要插入的数据
    data = [
        [chunk["page_num"] for chunk in chunks],  # 页码列表
        [chunk["chunk"] for chunk in chunks],  # 文本列表
        embeddings  # 向量列表
    ]

    # 插入数据
    mr = collection.insert(data)
    print(f"成功插入 {len(mr.primary_keys)} 条数据")

    # 加载集合到内存（必须操作，否则无法搜索）
    collection.load()


# 测试代码
if __name__ == "__main__":
    # 1. 初始化Milvus集合
    collection = init_milvus()

    # 2. 模拟数据（实际使用时从前面的步骤获取）
    sample_chunks = [
        {"page_num": 1, "chunk": "平安保险成立于1988年，总部位于深圳。"},
        {"page_num": 1, "chunk": "旗下有多个保险产品，包括寿险、车险、健康险等。"},
        {"page_num": 2, "chunk": "平安福是公司的明星寿险产品，2024年推出了新版。"}
    ]

    # 模拟向量（实际使用时调用text_embeddings函数生成）
    # 这里用随机数代替真实向量（仅测试用）
    sample_embeddings = [np.random.rand(1024).tolist() for _ in range(3)]

    # 3. 插入数据到Milvus
    insert_into_milvus(collection, sample_chunks, sample_embeddings)
